import torch
import cv2
import matplotlib.pyplot as plot

def test():

    image_rgb = load_image() # 加载图像
    results = act(image_rgb) # 进行推理
    detections = results.pandas().xyxy[0]  #解析结果: 转换为Pandas DataFrame

    # 打印检测结果
    print("检测到的对象:")
    print(detections[['name', 'confidence', 'xmin', 'ymin', 'xmax', 'ymax']])

    # 可视化结果
    results.render()  # 在图像上绘制边界框和标签
    output_image = results.ims[0]  # 获取带标注的图像
    display(output_image)

    # 保存结果
    output_path = 'detection_result.jpg'
    cv2.imwrite(output_path, cv2.cvtColor(output_image, cv2.COLOR_RGB2BGR))
    print(f"结果已保存至: {output_path}")


def act(image_rgb):
    model = torch.hub.load('ultralytics/yolov5', 'custom', path='../files/yolov5s.pt')
    device = 'cuda' if torch.cuda.is_available() else 'mps'
    model.to(device)
    results = model(image_rgb)  # 进行推理
    return results


def load_image():
    # 图像路径（替换为你自己的图片路径）
    # image_path = '../files/zidane.jpg'  # 可识别人的轮廓
    # image_path = '../files/fire_woods.webp'  # 无法识别
    image_path = '../files/old01.webp'  # 可识别人的轮廓
    image = cv2.imread(image_path)
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 转换为RGB格式
    return image_rgb


def display(output_image):
    # 显示结果
    plot.figure(figsize=(12, 8))
    plot.imshow(output_image)
    plot.axis('off')
    plot.title('YOLOv5 目标检测结果')
    plot.show()


if __name__ == '__main__':
    test()